#include "rclcpp/rclcpp.hpp"
#include "sensor_msgs/msg/image.hpp"
#include "nx_messenger/msg/box.hpp"
#include "nx_processor/yolo.h"

void setParameters(utils::InitParameter& initParameters)
{
	initParameters.class_names = utils::dataSets::armor;
	initParameters.num_class = 2; 
	initParameters.batch_size = 4;//undetermined
	initParameters.dst_h = 640;
	initParameters.dst_w = 640;
	initParameters.input_output_names = { "images",  "output0" };
	initParameters.conf_thresh = 0.25f;
	initParameters.iou_thresh = 0.45f;
	
}

void task(YOLOV8& yolo, const utils::InitParameter& param, std::vector<cv::Mat>& imgsBatch, const int& delayTime, const int& batchi,
	const bool& isShow)
{
	utils::DeviceTimer d_t0; yolo.copy(imgsBatch);	      float t0 = d_t0.getUsedTime();
	utils::DeviceTimer d_t1; yolo.preprocess(imgsBatch);  float t1 = d_t1.getUsedTime();
	utils::DeviceTimer d_t2; yolo.infer();				  float t2 = d_t2.getUsedTime();
	utils::DeviceTimer d_t3; yolo.postprocess(imgsBatch); float t3 = d_t3.getUsedTime();
	sample::gLogInfo << 
		"copy time = " << t0 / param.batch_size << "; "
		"preprocess time = " << t1 / param.batch_size << "; "
		"infer time = " << t2 / param.batch_size << "; "
		"postprocess time = " << t3 / param.batch_size << std::endl;

	if(isShow)
		utils::show(yolo.getObjectss(), param.class_names, delayTime, imgsBatch);
	
	yolo.reset();
}

class MessengerNode : public rclcpp::Node
{
public:
    MessengerNode()
        : Node("nx_messenger")
    {
        publisher_ = this->create_publisher<nx_messenger::msg::Box>("/nx_messenger/box", 10);
        subscription_ = this->create_subscription<sensor_msgs::msg::Image>(
            "/camera/color/image_raw", 10, std::bind(&MessengerNode::image_callback, this, std::placeholders::_1));
    }

private:
    void image_callback(const sensor_msgs::msg::Image::SharedPtr msg)
    {
        utils::DeviceTimer d_t0; yolo.copy(imgsBatch);	      float t0 = d_t0.getUsedTime();
	    utils::DeviceTimer d_t1; yolo.preprocess(imgsBatch);  float t1 = d_t1.getUsedTime();
	    utils::DeviceTimer d_t2; yolo.infer();				  float t2 = d_t2.getUsedTime();
        utils::DeviceTimer d_t3; yolo.postprocess(imgsBatch); float t3 = d_t3.getUsedTime();
        sample::gLogInfo << 
		"copy time = " << t0 / param.batch_size << "; "
		"preprocess time = " << t1 / param.batch_size << "; "
		"infer time = " << t2 / param.batch_size << "; "
		"postprocess time = " << t3 / param.batch_size << std::endl;

		objects = yolo.getObjectss();
		
        // TODO: Call your neural network model here and get the bounding box

        nx_messenger::msg::Box box_msg;
        box_msg.global_w = param.dst_w;
        box_msg.global_h = param.dst_h;
        box_msg.left_up_x = objects.left;
        box_msg.left_up_y = objects.top;
        cv::Mat img;
        box_msg.image = cv::rectangle(img, cv::Point(objects.left, objects.top), cv::Point(objects.right, objects.bottom), cv::Scalar(0, 0, 255), 2); 
        // TODO: Fill the box_msg fields here

        publisher_->publish(box_msg);
        yolo.reset();
    }

    rclcpp::Publisher<nx_messenger::msg::Box>::SharedPtr publisher_;
    rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr subscription_;
};

int main(int argc, char *argv[])
{

    // parameters
    utils::InitParameter param;
	setParameters(param);

	// path
	std::string model_path = "../best.trt";

	YOLOV8 yolo(param);

	// read model
	std::vector<unsigned char> trt_file = utils::loadModel(model_path);
	if (trt_file.empty())
	{
		sample::gLogError << "trt_file is empty!" << std::endl;
		return -1;
	}
	// init model
	if (!yolo.init(trt_file))
	{
		sample::gLogError << "initEngine() ocur errors!" << std::endl;
		return -1;
	}
	yolo.check();
	cv::Mat frame;
	std::vector<cv::Mat> imgs_batch;
	imgs_batch.reserve(param.batch_size);
	sample::gLogInfo << imgs_batch.capacity() << std::endl;


    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<MessengerNode>());
    rclcpp::shutdown();
    return 0;
}